#ifndef DIFF_AVL_H_
#define DIFF_AVL_H_
#include "avl_tree.h"
#include "vector.h"

class DiffAVL : public AVLTree<Vector, false> {
public:
	bool insert(const Vector& val) {
		bool res;
		if (this->root == 0){	//first node
			this->root = new AVLNode<Vector>(val);
			return true;
		}
		int prev;
		AVLNode<Vector>* curr = search(val, res, prev);
		if (res) return false;
		AVLNode<Vector>* newNode = new AVLNode<Vector>(val);
		//keep the difference from parent instead of actual value
		int new_a = get_diff(newNode) - prev;
		Vector vec(new_a, newNode->value().get_b());
		newNode->setValue(vec);
		newNode->connect(curr, PARENT);
		vec.set_a(prev);
		vec.set_b(curr->value().get_b());
		if ( val < vec) 
			curr->connect(newNode, LEFT);
		else
			curr->connect(newNode, RIGHT);

		Balance(newNode, val.get_a());
		return true;	
	}
	
	bool remove(const Vector& val){
		bool res, oneChild = false, twoChildren = false;
		NodeType nt;
		int curr_val;
		AVLNode<Vector>* curr = search(val, res, curr_val);
		int parent_val;
		if (!res) return false;
		
		if (curr->getTree(LEFT) && curr->getTree(RIGHT)){ //if curr has 2 children
			curr = twoChildrenRemove(curr, curr_val);
			twoChildren = true;
		}
		
		//check if the has only one child
		if (curr->getTree(LEFT)){
			if (!curr->getTree(RIGHT)) {
				nt = LEFT;
				oneChild = true;
			}
		} else{
			if (curr->getTree(RIGHT)){
				nt = RIGHT;
				oneChild = true;
			}
		}	
	
		AVLNode<Vector>* parent = curr->getTree(PARENT);
		parent_val = curr_val - get_diff(curr);
		
		if (oneChild) (curr->getTree(nt))->connect(parent, PARENT);
		
		if (curr->isLeaf() || oneChild){
			if (root == curr){ 				
				root = oneChild ? root->getTree(nt) : 0;
				if (oneChild){
					Vector v(curr_val + get_diff(curr->getTree(nt)), 
						curr->getTree(nt)->value().get_b());
					root->setValue(v);
				}
				delete curr;
			} else {
				NodeType nt2;
				nt2 = ( parent->getTree(LEFT) == curr ) ? 
							LEFT : RIGHT;
				parent->disconnect(nt2);
				if (oneChild){ 
					Vector v(get_diff(curr->getTree(nt)) + get_diff(curr), 
							curr->getTree(nt)->value().get_b());
					curr->getTree(nt)->setValue(v);
					(parent)->connect(curr->getTree(nt), nt2);
				}

				delete curr;
			}
		}
		
		Balance(parent, parent_val);
		return true;
	}

	bool exist (const Vector& val) const{ 
		bool b;
		int prev;
		search(val, b, prev);
		return b;
	}
	
	int rank(const Vector& val) const{
		if (!root) return -1;
		AVLNode<Vector>* curr = root;
		Vector v(0,0);
		int tot = 0;

		while (true) {
			v.set_a(v.get_a() + get_diff(curr));
			v.set_b(curr->value().get_b());
			
			if (val == v){
				tot += curr->getLeftChildren() + 1;
				break;
			}
			
			if (curr->isLeaf()) {
				if (!(val < v)) tot += curr->getLeftChildren() + 1;
				break;
			}
			
			if (val < v) {
				if  (curr->getTree(LEFT) == 0) {
					break;
				}
				curr = curr->getTree(LEFT);
			} else {
				if  (curr->getTree(RIGHT) == 0) {
					tot += curr->getLeftChildren() + 1;
					break;
				}
				tot += curr->getLeftChildren() + 1;
				curr = curr->getTree(RIGHT);
			}
		}
		return tot;
				
	}
	
	bool update(const int a, const int k) {
		AVLNode<Vector>* curr = root;
		Vector v(0,0);
		bool stop = false, go_right = true;
		int val;
		if (root == 0) {
			return false;
		}
		
		val = get_diff(root);
		
		while (!stop){
			if (val > a && go_right){			//continue right (if possible)
				if (curr->getTree(RIGHT)){
					curr = curr->getTree(RIGHT);
					val += get_diff(curr);	
				} else
					stop = true;
			} else if (val <= a && !go_right){	//continue left (if possible)
				if (curr->getTree(LEFT)){
					curr = curr->getTree(LEFT);
					val += get_diff(curr);	
				} else
					stop = true;		
			} else if (val <= a && go_right){	//update curr & continue left
				if (curr != root) val += get_diff(curr);	
				v.set_a(get_diff(curr) - k);
				v.set_b(curr->value().get_b());
				curr->setValue(v);		
				go_right = false;
			} else if (val > a && !go_right){	//fix curr & continue right
				val += get_diff(curr);	
				v.set_a(get_diff(curr) + k);
				v.set_b(curr->value().get_b());
				curr->setValue(v);		
				go_right = true;

			}
		}		
	
		return true;
	}
	
private: 
	
	static inline int get_diff(const AVLNode<Vector>* node) {return  (node == 0)? 0 : node->value().get_a();}

	AVLNode<Vector>* search(const Vector& val, bool& result, int& sum) const{
		AVLNode<Vector>* curr = root;
		Vector v(0, 0);
		sum = 0;
		if (root == 0) {
			result = false;
			return 0;
		}
		
		while (true) {
			sum += get_diff(curr);
			v.set_a(v.get_a() + get_diff(curr));
			v.set_b(curr->value().get_b());
			
			if (val == v){
				result = true;
				return curr;
			}
			
			if (curr->isLeaf()) {
				result = false;
				return curr;
			}
			
			if (val < v) {
				if  (curr->getTree(LEFT) == 0) {
					result = false;
					return curr;
				}
				curr = curr->getTree(LEFT);
			} else {
				if  (curr->getTree(RIGHT) == 0) {
					result = false;
					return curr;
				}
				curr = curr->getTree(RIGHT);
			}
		}
	}

	static void connectParent(AVLNode<Vector>* node, int val) {
		AVLNode<Vector>* parent = node->getTree(PARENT);
		if (parent == 0) return;
		Vector v_p(val - get_diff(node) , parent->value().get_b()), v_n(val, node->value().get_b());

		if (v_n < v_p)
			parent->connect(node, LEFT);
		else
			parent->connect(node, RIGHT);
		/*
		if (val > (val - get_diff(node)))
			parent->connect(node, LEFT);
		else if (val < (val - get_diff(node)))
			parent->connect(node, RIGHT);
		else{
			if (node->value().get_b() >= parent->value().get_b()){
				parent->connect(node, LEFT);
			} else{
				parent->connect(node, RIGHT);
			}
		}*/
	}

	static AVLNode<Vector>* getPrevInOrder(AVLNode<Vector>* node, int& val){
		AVLNode<Vector>* temp = node;
		temp = temp->getTree(LEFT);
		val += temp->value().get_a();
		while (temp->getTree(RIGHT)) {
			temp = temp->getTree(RIGHT);
			val += temp->value().get_a();
		}
		return temp;
	}
	
	//Switching the node to remove with the prev in order, returns the second & it's value.
	static AVLNode<Vector>* twoChildrenRemove(AVLNode<Vector>* curr, int& curr_val){			
			int temp_val = curr_val;
			AVLNode<Vector>* temp = getPrevInOrder(curr, temp_val);
			
			// set the curr & it's children
			int parent_val = curr_val - get_diff(curr);
			int l_val = curr_val + get_diff(curr->getTree(LEFT));
			int r_val = curr_val + get_diff(curr->getTree(RIGHT));
			int curr_b = curr->value().get_b();
			Vector v(temp_val - parent_val, temp->value().get_b());
			curr->setValue(v);
			v.set_a(l_val - temp_val);
			v.set_b(curr->getTree(LEFT)->value().get_b());
			curr->getTree(LEFT)->setValue(v);
			v.set_a(r_val - temp_val);
			v.set_b(curr->getTree(RIGHT)->value().get_b());
			curr->getTree(RIGHT)->setValue(v);
			
			//set the temp & it's children
			parent_val = temp_val - get_diff(temp);
			l_val = temp_val + get_diff(curr->getTree(LEFT));
			r_val = temp_val + get_diff(curr->getTree(RIGHT));
			v.set_a(curr_val - parent_val);
			v.set_b(curr_b);
			temp->setValue(v);
			if (temp->getTree(LEFT)){
				v.set_a(l_val - curr_val);			
				v.set_b(temp->getTree(LEFT)->value().get_b());
				temp->getTree(LEFT)->setValue(v);
			}
			if (temp->getTree(RIGHT)){
				v.set_a(r_val - curr_val);			
				v.set_b(temp->getTree(RIGHT)->value().get_b());
				temp->getTree(RIGHT)->setValue(v);	
			}	
			
			return temp;
	}

	//Balancing the tree after inserting the given node or deleting the node's child
	void Balance(AVLNode<Vector>* node, int val){
		AVLNode<Vector>* parent;
		AVLNode<Vector>* grand;
		AVLNode<Vector>* temp;
		int parent_val;
		int temp_val;
		
		if (!node  || !root) return;
		
		// we first need to check the *current* node (not the parent)
		if (node->getTree(RIGHT)){
			node = node->getTree(RIGHT);
			val += get_diff(node);
		} else if (node->getTree(LEFT)){
			node = node->getTree(LEFT);
			val += get_diff(node);
		}
		
		while (node != root) {
			parent = node->getTree(PARENT);
			parent_val = parent ? val - get_diff(node) : 0;
			grand = parent->getTree(PARENT);
			RotType rot = getRotType(parent);
			switch (rot){
				case LL:	RotateLeft(parent, parent_val);
							break;
				case LR:	temp = parent->getTree(LEFT);
							temp_val = parent_val + get_diff(temp);
							RotateRight(temp, temp_val);
							RotateLeft(parent, parent_val);
							break;
				case RR:	RotateRight(parent, parent_val);
							break;
				case RL:	temp = parent->getTree(RIGHT);
							temp_val = parent_val + get_diff(temp);
							RotateLeft(temp, temp_val);
							RotateRight(parent, parent_val);
				default:	break;
			}
			
			if (rot == NOROT) connectParent(node, val);

			if (rot != NOROT && !grand) 
				root = parent;
			else 
				connectParent(node, val);
	
			node = parent;
			val = parent_val;
		}
	}

	static void RotateLeft(AVLNode<Vector>*& root, int& val) {
		AVLNode<Vector>* grand = root->getTree(PARENT);
		AVLNode<Vector>* l = root->getTree(LEFT);
		AVLNode<Vector>* lr = l->getTree(RIGHT); // check corrrectness
		
		int grand_val = (grand == 0) ? 0 : val - get_diff(root);
		int l_val = val + get_diff(l);
		int lr_val = l_val + get_diff(lr);
		
		Vector temp(l_val - grand_val , l->value().get_b()); //work also if root is tree root
		l->setValue(temp);
		
		temp.set_a(val - l_val);
		temp.set_b(root->value().get_b());
		root->setValue(temp);		
		
		if (lr){
			temp.set_a(lr_val - val);
			temp.set_b(lr->value().get_b());
			lr->setValue(temp);
		} 
		
		
		if (lr) lr->connect(root, PARENT);
		root->connect(root->getTree(LEFT), PARENT);
		root->connect(lr, LEFT);
		(root->getTree(PARENT))->connect(root, RIGHT);
		root = root->getTree(PARENT);
		root->connect(grand, PARENT);
		connectParent(root, l_val);
		val = l_val;
	}
	
	static void RotateRight (AVLNode<Vector>*& root, int& val) {
		AVLNode<Vector>* grand = root->getTree(PARENT);
		AVLNode<Vector>* r = root->getTree(RIGHT);
		AVLNode<Vector>* rl = r->getTree(LEFT); // check corrrectness

		int grand_val = (grand == 0) ? 0 : val - get_diff(root);
		int r_val = val + get_diff(r);
		int rl_val = r_val + get_diff(rl);
		
		Vector temp(r_val - grand_val , r->value().get_b()); //work also if root is tree root
		r->setValue(temp);
		
		temp.set_a(val - r_val);
		temp.set_b(root->value().get_b());
		root->setValue(temp);	

		if (rl){
			temp.set_a(rl_val - val);
			temp.set_b(rl->value().get_b());
			rl->setValue(temp);
		}		
		

		if (rl) rl->connect(root, PARENT);
		root->connect(root->getTree(RIGHT), PARENT);
		root->connect(rl, RIGHT);
		(root->getTree(PARENT))->connect(root, LEFT);
		root = root->getTree(PARENT);
		root->connect(grand, PARENT);
		connectParent(root, r_val);
		val = r_val;
	}
};
#endif


/***
 TODO:
 1) remove
 2) rotations
 3) update
 */
